import importlib.util
import importlib
import json
import socket
from socket import *
from threading import *
import json
import argparse

parser = argparse.ArgumentParser(description="wargame client program")
parser.add_argument("-f", "--faction", type=int, default=0, help="所属方，0:红方，1:蓝方")
parser.add_argument('-n', '--user_name', type=str, default='demo', help='用户名')
parser.add_argument('-i', '--user_id', type=int, default=0, help='用户编号')
parser.add_argument('--agent_file_path', type=str, default='ai/agent.py', help='AI文件所在（相对）路径')
parser.add_argument('--agent_class_name', type=str, default="Agent", help='AI类名')
parser.add_argument('-p', '--port', type=int, default=7111, help='端口号：7111或7211')
args = parser.parse_args()


def sendMsg(sock,msg,addr):
    #while True:
    n=sock.sendto(msg.encode(), addr)


def recvMsg(sock, addr, myai):
    # sock.sendto('client_name：'.encode(), addr)
    while True:
        data = sock.recvfrom(89120)
        if "ok" in data[0].decode('gb18030')[:2]:
            print(data[0].decode('gb18030'))
            # sendMsg(sock, "begin" ,addr)  # 2023-11-21 不再需要发送begin消息
            # if initial == False:#客户端收到态势后更新动作，并发送给服务器端
            #   buffer = myai.step(myai.observation)
            #   sendMsg(sock,buffer,addr)
        else:
            observation = {}
            observation_string=""
            recv_str = data[0].decode( encoding = 'gb2312')
            # recv_str=replace_string(recv_str,"[{","{{")
            # recv_str = replace_string(recv_str, "}]","}}")
            # for key_value in recv_str.split(','):
            #     [key, value] = key_value.strip().split(":")[:2]
            #     observation[key] = value
            observation=json.loads(recv_str)
            modify_observation(observation)
            # for key, value in observation.items():
            #     if isinstance(value, list):
            #         observation[key] = value
            #     elif isinstance(value, int):
            #         observation[key] = [value]

            print(f"*** side: {'红' if myai.color == 0 else '蓝'} *** cur_step: {observation['time']['cur_step']} ***")
            # print("--- valid_actions ---")  # 调试用
            # print(observation["valid_actions"])
            dic_obser=myai.step(observation)
            observation_string = json.dumps(dic_obser)
            action ="act"+ observation_string
            sendMsg(sock,action,addr)
            print("--- actions ---")
            print(action)  # 调试用
            #initial =False


def replace_string(original_string, target_string, replacement_string):
    # 使用replace()函数替换指定字符串
    new_string = original_string.replace(target_string, replacement_string)
    return new_string


def modify_observation(observation):
    if 'operators' in observation.keys():
        for opr in observation['operators']:
            observe_distance = {}
            if opr['observe_distance']:
                for item in opr['observe_distance']:
                    observe_distance[item['Key']] = item['Value']
                opr['observe_distance'] = observe_distance

            remain_bullet_nums = {}
            if opr['remain_bullet_nums']:
                for item in opr['remain_bullet_nums']:
                    remain_bullet_nums[item['Key']] = item['Value']
                opr['remain_bullet_nums'] = remain_bullet_nums

            remain_bullet_nums_bk = {}
            if opr['remain_bullet_nums_bk']:
                for item in opr['remain_bullet_nums_bk']:
                    remain_bullet_nums_bk[item['Key']] = item['Value']
                opr['remain_bullet_nums_bk'] = remain_bullet_nums_bk
    if 'passengers' in observation.keys():
        if observation['passengers']:
            for opr in observation['passengers']:
                observe_distance = {}
                if opr['observe_distance']:
                    for item in opr['observe_distance']:
                        observe_distance[item['Key']] = item['Value']
                    opr['observe_distance'] = observe_distance

            remain_bullet_nums = {}
            if opr['remain_bullet_nums']:
                for item in opr['remain_bullet_nums']:
                    remain_bullet_nums[item['Key']] = item['Value']
                opr['remain_bullet_nums'] = remain_bullet_nums

            remain_bullet_nums_bk = {}
            if opr['remain_bullet_nums_bk']:
                for item in opr['remain_bullet_nums_bk']:
                    remain_bullet_nums_bk[item['Key']] = item['Value']
                opr['remain_bullet_nums_bk'] = remain_bullet_nums_bk
    if 'valid_actions' in observation.keys():
        valid_actions = {}
        if observation['valid_actions']:
            for action in observation['valid_actions']:
                action_value = {}
                for val in action['Value']:
                    action_value[val['Key']] = val['Value']
                valid_actions[action['Key']] = action_value
        observation['valid_actions'] =  valid_actions


def create_instance_from_file(file_path, class_name):
    spec = importlib.util.spec_from_file_location(class_name, file_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return getattr(module, class_name)()


def main():
    # 自定义初始化输入
    faction = args.faction  # 红方：0，蓝方：1
    seat = faction  # 席位（目前未使用）
    if faction == 0: color = "red" 
    else: color = "blue"
    role = faction  # 角色（目前未使用）
    user_name = args.user_name  # 自定义名称
    user_id = args.user_id  # 用户ID
    map_id = 94  # 地图编号

    scenario_id = 2010131194  # 使用的想定
    agent_file_path = args.agent_file_path  # AI文件路径
    agent_class_name = args.agent_class_name  # AI的类名
    
    ip_address = "127.0.0.1"  # 服务器IP地址
    port = args.port  # 服务器端口

    # 第一步：服务器发送开始推演指令，客户端实例化agent
    myai = create_instance_from_file(agent_file_path, agent_class_name)
    
    with open(f"data/scenarios/{scenario_id}.json",encoding='utf-8') as f:
        scenario_data = json.load(f)
    myai.setup({
            "scenario": scenario_data,
            "map_id": map_id,
            "seat": seat,
            "faction": faction,
            "role": role,
            "user_name": user_name,
            "user_id": user_id,
        })
    
    # 第二步,连接服务器
    s = socket(AF_INET, SOCK_DGRAM)
    server = (ip_address, port)
    if faction == 0:
        client_port = 12340
    else:
        client_port = 12341
    s.bind(("",client_port))
    # msg=""
    sendMsg(s, color, server)
    tr = Thread(target=recvMsg, args=(s,server,myai), daemon=True)
    #ts = Thread(target=sendMsg, args=(s, msg,server))
    tr.start()
    # ts.start()
    while tr.is_alive():
        tr.join(1)


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("KeyboardInterrupt")
        exit()
